import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
from tensorflow.keras import backend as K

def cbam(inputs):
    outputs = channel_attention(inputs)
    outputs = spatial_attention(outputs)
    return outputs

def channel_attention(inputs):
    channel = inputs.shape[-1]
    
    avg_pool = layers.GlobalAveragePooling2D()(inputs)
    avg_pool = layers.Reshape((1, 1, channel))(avg_pool)
    avg_pool = layers.Dense(channel//8, activation="relu")(avg_pool)
    avg_pool = layers.Dense(channel, activation="relu")(avg_pool)
    
    max_pool = layers.GlobalMaxPooling2D()(inputs)
    max_pool = layers.Reshape((1, 1, channel))(max_pool)
    max_pool = layers.Dense(channel//8, activation="relu")(max_pool)
    max_pool = layers.Dense(channel, activation="relu")(max_pool)

    outputs = layers.Add()([avg_pool, max_pool])
    outputs = layers.Activation("sigmoid")(outputs)
    return layers.Multiply()([inputs, outputs])
    

def spatial_attention(inputs):
    avg_pool = layers.Permute((3, 2, 1))(inputs)
    avg_pool = layers.AveragePooling2D(pool_size=(inputs.shape[-1], 1))(avg_pool)
    avg_pool = layers.Permute((3, 2, 1))(avg_pool)
    max_pool = layers.Permute((3, 2, 1))(inputs)
    max_pool = layers.MaxPooling2D(pool_size=(inputs.shape[-1], 1))(max_pool)
    max_pool = layers.Permute((3, 2, 1))(max_pool)
    #avg_pool = layers.Lambda(lambda x: K.mean(x, axis=3, keepdims=True))(inputs)
    #max_pool = layers.Lambda(lambda x: K.max(x, axis=3, keepdims=True))(inputs)
    outputs = layers.Concatenate(axis=3)([avg_pool, max_pool])
    outputs = layers.Conv2D(filters=1, kernel_size=3, padding="same", activation="sigmoid")(outputs)
    return layers.Multiply()([inputs, outputs])

def gru_attention_block(inputs, length=15):
    outputs = gru_channel(inputs)
    outputs = gru_spatial(outputs)
    return outputs

def gru_channel(inputs):
    channel = inputs.shape[-1]
    
    avg_pool = layers.GlobalAveragePooling1D()(inputs)
    avg_pool = layers.Reshape((1, channel))(avg_pool)
    avg_pool = layers.Dense(channel//8, activation="relu")(avg_pool)
    avg_pool = layers.Dense(channel, activation="relu")(avg_pool)

    max_pool = layers.GlobalMaxPooling1D()(inputs)
    max_pool = layers.Reshape((1, channel))(max_pool)
    max_pool = layers.Dense(channel//8, activation="relu")(max_pool)
    max_pool = layers.Dense(channel, activation="relu")(max_pool)

    outputs = layers.Add()([avg_pool, max_pool])
    outputs = layers.Activation("sigmoid")(outputs)

    return layers.Multiply()([inputs, outputs])

def gru_spatial(inputs):
    avg_pool = layers.Permute((2, 1))(inputs)
    avg_pool = layers.AveragePooling1D(pool_size=inputs.shape[-1])(avg_pool)#错误点need 1 received 3
    avg_pool = layers.Permute((2, 1))(avg_pool)

    max_pool = layers.Permute((2, 1))(inputs)
    max_pool = layers.MaxPooling1D(pool_size=inputs.shape[-1])(max_pool)
    max_pool = layers.Permute((2, 1))(max_pool)

    outputs = layers.Concatenate(axis=-1)([avg_pool, max_pool])
    outputs = layers.Conv1D(filters=1, kernel_size=3, padding="same", activation="sigmoid")(outputs)
    return layers.Multiply()([inputs, outputs])
